# utils/bleu4.py
import nltk
from nltk.translate.bleu_score import sentence_bleu, SmoothingFunction
import numpy as np

def calculate_bleu4(reference, hypothesis):
    """计算BLEU-4分数
    
    Args:
        reference: 参考翻译（列表）
        hypothesis: 模型生成的翻译（列表）
        
    Returns:
        float: BLEU-4分数
    """
    # 检查输入是否为空
    if not reference or not hypothesis:
        return 0.0
    
    # 过滤掉特殊token
    reference = [token for token in reference if token not in ['<pad>', '<sos>', '<eos>', '<unk>']]
    hypothesis = [token for token in hypothesis if token not in ['<pad>', '<sos>', '<eos>', '<unk>']]
    
    # 如果过滤后序列为空，返回0
    if not reference or not hypothesis:
        return 0.0
    
    # 使用平滑函数
    smoothing = SmoothingFunction().method4
    
    try:
        # 计算BLEU-4分数
        weights = (0.25, 0.25, 0.25, 0.25)  # 4-gram权重
        bleu = sentence_bleu([reference], hypothesis, weights=weights, smoothing_function=smoothing)
        return bleu
    except Exception as e:
        print(f"Error calculating BLEU: {e}")
        print(f"Reference: {reference}")
        print(f"Hypothesis: {hypothesis}")
        return 0.0